import os
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
import random


class PredictDataset(BaseDataset):
    def __init__(self, opt):
        """Initialize this dataset class.
        初始化一个预测类，单个图片为feature
        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        self.dir_A = os.path.join(opt.dataroot, opt.phase)

        self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))


        self.A_size = len(self.A_paths)  # get the size of dataset A
        btoA = self.opt.direction == 'BtoA'
        input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
        
        self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))

    def __getitem__(self, index):
        A_path = self.A_paths[index % self.A_size]  # make sure index is within then range
        A_img = Image.open(A_path).convert('RGB')
        A = self.transform_A(A_img)

        return {'A': A, 'A_paths': A_path}

    def __len__(self):
        return self.A_size
